package logger

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"io"
	"time"
)

var maskConfig = MaskFieldsConfig{
	FieldsToMask: []string{"password", "token"},
	MaskedValue:  "*****",
}

type MaskFieldsConfig struct {
	FieldsToMask []string
	MaskedValue  string
}

type MiddlewareBuilder struct {
	allowReqBody  bool
	allowRespBody bool
	loggerFunc    func(ctx context.Context, al *AccessLog)
}

func NewBuilder(fn func(ctx context.Context, al *AccessLog)) *MiddlewareBuilder {
	return &MiddlewareBuilder{
		loggerFunc: fn,
	}
}

func (b *MiddlewareBuilder) AllowReqBody() *MiddlewareBuilder {
	b.allowReqBody = true
	return b
}

func (b *MiddlewareBuilder) AllowRespBody() *MiddlewareBuilder {
	b.allowRespBody = true
	return b
}

func (b *MiddlewareBuilder) Build() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		start := time.Now()
		url := ctx.Request.URL.String()
		if len(url) > 1024 {
			url = url[:1024]
		}
		al := &AccessLog{
			Method: ctx.Request.Method,
			Status: ctx.Writer.Status(),
			Url:    url,
		}
		if b.allowReqBody && ctx.Request.Body != nil {
			body, _ := ctx.GetRawData()
			ctx.Request.Body = io.NopCloser(bytes.NewReader(body))

			if len(body) > 1024 {
				body = body[:1024]
			}
			al.ReqBody = string(body)
		}

		if b.allowRespBody {
			ctx.Writer = responseWriter{
				al:             al,
				ResponseWriter: ctx.Writer,
			}
		}

		defer func() {
			al.Duration = time.Since(start).String()
			b.loggerFunc(ctx, al)
		}()
		ctx.Next()
	}
}

type responseWriter struct {
	al *AccessLog
	gin.ResponseWriter
}

func (w responseWriter) WriteHeader(statusCode int) {
	w.al.Status = statusCode
	w.ResponseWriter.WriteHeader(statusCode)
}

func (w responseWriter) Write(data []byte) (int, error) {
	w.al.RespBody = string(data)
	return w.ResponseWriter.Write(data)
}

func (w responseWriter) WriteString(data string) (int, error) {
	w.al.RespBody = data
	return w.ResponseWriter.WriteString(data)
}

type AccessLog struct {
	// HTTP 请求的方法
	Method string
	// Url 整个请求 URL
	Url      string
	Duration string
	ReqBody  string
	RespBody string
	Status   int
}

func (al AccessLog) String() string {
	modifiedReqBody := al.modifyFields(al.ReqBody)
	return fmt.Sprintf("Method:%s | Url:%s | Duration:%s | Status:%d | ReqBody:%s ", al.Method, al.Url, al.Duration, al.Status, modifiedReqBody)
}
func (al AccessLog) modifyFields(inputJSON string) string {
	var jsonMap map[string]interface{}
	if err := json.Unmarshal([]byte(inputJSON), &jsonMap); err != nil {
		return inputJSON
	}

	for _, field := range maskConfig.FieldsToMask {
		if _, ok := jsonMap[field]; ok {
			jsonMap[field] = maskConfig.MaskedValue
		}
	}

	modifiedJSONBytes, err := json.Marshal(jsonMap)
	if err != nil {
		return inputJSON
	}

	return string(modifiedJSONBytes)
}
